package com.itheima.ai.controller;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.io.FileTypeUtil;
import com.itheima.ai.repository.FileRepository;
import com.itheima.ai.vo.Result;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.ExtractedTextFormatter;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.config.PdfDocumentReaderConfig;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.Resource;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.List;

/**
 * RAG知识库-PDF文件处理控制器
 *
 * @Author mr.wu
 * @Date 2025-11-6 16:40
 */
@CrossOrigin
@RestController
@RequestMapping("/ai/pdf")
public class FileController {

    @Autowired
    private FileRepository fileRepository;


    /**
     * PDF 文件上传接口
     * @param chatId  对话ID
     * @param file  PDF文件对象
     * @return
     */
    @PostMapping("/upload/{chatId}")
    public Result upload(@PathVariable String chatId, MultipartFile file){
        try {
            //1. 判断是否PDF文件
            String type = FileTypeUtil.getType(file.getInputStream());
            if(file.isEmpty() || file.getSize()==0 || !type.contains("pdf")) {
                return Result.fail("非法文件");
            }

            //2. 保存PDF到磁盘和与会话ID的关系
            fileRepository.savePdfFile(chatId,file.getResource());

            //3. 将PDF内容写入向量数据库
            this.writePdf2Vectore(file.getResource());

            return Result.ok();
        } catch (IOException e) {
            e.printStackTrace();
            return Result.fail("上传失败");
        }
    }

    /**
     * PDF下载接口（方便前端预览）
     * @param chatId
     * @return
     */
    @GetMapping("/file/{chatId}")
    public ResponseEntity download(@PathVariable String chatId) {
        //1. 获取文件资源
        Resource resource = fileRepository.getPdfFile(chatId);

        //2. 判断是否存在
        if(resource==null) {
            return ResponseEntity.notFound().build();
        }

        //3. 处理下载
        String filename = resource.getFilename();
        filename= URLEncoder.encode(filename, StandardCharsets.UTF_8);
        return ResponseEntity.ok() //200状态码
                .contentType(MediaType.APPLICATION_OCTET_STREAM) //固定文件类型
                .header(HttpHeaders.CONTENT_DISPOSITION,"attachment;filename="+filename) //固定的响应体
                .body(resource); //要下载的文件资源
    }

    @Autowired
    private VectorStore vectorStore;

    private void writePdf2Vectore(Resource resource) {
        //1. 创建PDF阅读器
        PagePdfDocumentReader reader = new PagePdfDocumentReader(resource, PdfDocumentReaderConfig.builder()
                .withPageExtractedTextFormatter(ExtractedTextFormatter.defaults())
                .withPagesPerDocument(1)//每页当做一个文档对象
                .build());

        //2. 读取文档集合
        List<Document> documentList = reader.read();

        //3. 写入向量数据库
        if(CollUtil.isNotEmpty(documentList)) {
            vectorStore.add(documentList);
        }
    }
}
